
"""数据安全模块，提供数据加密、解密和访问控制功能"""
import hashlib
import json
import base64
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from typing import Dict, Any, Optional, List
from datetime import datetime, timedelta

class DataSecurityManager:
    """数据安全管理器"""
    
    def __init__(self, secret_key: Optional[str] = None):
        # 如果未提供密钥，使用默认密钥 (实际应用中应该从配置或环境变量中读取)
        if secret_key is None:
            secret_key = "default_fin_senti_secret_key"  # 实际应用中应该从安全的地方获取
        
        # 从密钥派生加密密钥
        salt = hashlib.sha256(secret_key.encode()).digest()[:16]
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(secret_key.encode()))
        
        self.cipher = Fernet(key)
        self.access_control_list = {}
    
    def encrypt_data(self, data: Any) -> str:
        """加密数据
        
        Args:
            data: 要加密的数据，可以是任何可JSON序列化的对象
            
        Returns:
            str: 加密后的数据
        """
        # 将数据序列化为JSON字符串
        data_str = json.dumps(data, ensure_ascii=False)
        
        # 加密数据
        encrypted_data = self.cipher.encrypt(data_str.encode('utf-8'))
        
        # 将加密后的数据转换为字符串
        return encrypted_data.decode('utf-8')
    
    def decrypt_data(self, encrypted_data: str) -> Any:
        """解密数据
        
        Args:
            encrypted_data: 加密后的数据
            
        Returns:
            Any: 解密后的数据
        """
        # 解密数据
        decrypted_data = self.cipher.decrypt(encrypted_data.encode('utf-8'))
        
        # 将解密后的数据反序列化为Python对象
        return json.loads(decrypted_data.decode('utf-8'))
    
    def hash_data(self, data: str, salt: Optional[str] = None) -> str:
        """对数据进行哈希处理
        
        Args:
            data: 要哈希的数据
            salt: 盐值
            
        Returns:
            str: 哈希后的数据
        """
        if salt:
            data = f"{data}{salt}"
        
        return hashlib.sha256(data.encode('utf-8')).hexdigest()
    
    def add_access_control(self, resource_id: str, allowed_roles: List[str], expiry: Optional[datetime] = None):
        """添加访问控制规则
        
        Args:
            resource_id: 资源ID
            allowed_roles: 允许访问的角色列表
            expiry: 过期时间
        """
        self.access_control_list[resource_id] = {
            'allowed_roles': allowed_roles,
            'expiry': expiry
        }
    
    def check_access(self, resource_id: str, user_roles: List[str]) -> bool:
        """检查用户是否有权限访问资源
        
        Args:
            resource_id: 资源ID
            user_roles: 用户角色列表
            
        Returns:
            bool: 是否有权限访问
        """
        # 检查资源是否存在访问控制规则
        if resource_id not in self.access_control_list:
            # 默认拒绝访问
            return False
        
        # 获取访问控制规则
        rule = self.access_control_list[resource_id]
        
        # 检查规则是否过期
        if rule['expiry'] and datetime.now() > rule['expiry']:
            return False
        
        # 检查用户角色是否在允许的角色列表中
        return any(role in rule['allowed_roles'] for role in user_roles)
    
    def mask_sensitive_info(self, data: Dict[str, Any], sensitive_fields: List[str]) -> Dict[str, Any]:
        """屏蔽敏感信息
        
        Args:
            data: 要处理的数据
            sensitive_fields: 敏感字段列表
            
        Returns:
            Dict: 处理后的数据
        """
        # 深拷贝数据以避免修改原始数据
        masked_data = json.loads(json.dumps(data))
        
        # 屏蔽敏感字段
        for field in sensitive_fields:
            if field in masked_data:
                value = masked_data[field]
                if isinstance(value, str):
                    if len(value) <= 4:
                        masked_data[field] = '*' * len(value)
                    else:
                        # 保留前两位和后两位
                        masked_data[field] = value[:2] + '*' * (len(value) - 4) + value[-2:]
            
        return masked_data